#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/5/30 17:52
# @Author  : huidong.bai
# @File    : audio_tag_split.py
# @Software: PyCharm
# @Mail    : MasterBai2018@outlook.com
import os
import sys
import numpy as np
import scipy.signal as signal
from pathlib import Path
import re


file_list = []


def extract_file_channel(file_path, total_channels, extract_channel):
    with open(file_path, 'rb') as f:
        if file_path.endswith(".wav"):
            # 获取文件头信息
            header = f.read(100)
            if header[:4] == b"RIFF":
                # 如果是wav文件，则获取头信息中的音频长度
                pos = re.search(b'data', header).end()
                audio_length = int.from_bytes(header[pos:pos + 4], byteorder='little')
                # 将文件指针移动到音频数据的起始位置
                f.seek(pos + 4, 0)
            else:
                f.seek(0, 0)
        audio_data = np.fromfile(f, dtype=np.int16)
        # 横向长度未知，纵向长度是通道数
        audio_data.shape = -1, total_channels
    audio_data = audio_data.T
    return audio_data[extract_channel-1]


def find_voice_tag(audio_file, total_channels, people_channel):
    sample_rate = 16000

    audio_data = extract_file_channel(audio_file, total_channels, people_channel)
    window_size = 1024
    hop_size = 512
    n_fft = 1024

    # Compute STFT
    f, t, Zxx = signal.stft(audio_data, fs=sample_rate, window='hann', nperseg=window_size, noverlap=hop_size,
                            nfft=n_fft)

    # Compute magnitude spectrogram
    mag_spec = np.abs(Zxx)
    # Compute frequency bins
    freq_bins = np.fft.fftfreq(n_fft, 1 / sample_rate)[:n_fft // 2]

    # Find index of frequency bin closest to 1kHz
    idx_1kHz = np.argmin(np.abs(freq_bins - 1000))
    idx_2kHz = np.argmin(np.abs(freq_bins - 1906))

    # Find index of frequency bin closest to 3kHz
    idx_3kHz = np.argmin(np.abs(freq_bins - 2906))

    # Find index of frequency bin closest to 5kHz
    idx_5kHz = np.argmin(np.abs(freq_bins - 4810))

    # Compute dB values for each frequency bin
    dB_spec = 20 * np.log10(mag_spec)
    max_indices = np.argmax(dB_spec, axis=0)
    indices = np.where(max_indices == idx_1kHz)[0]

    di_tag = -1
    for x in np.nditer(indices):
        if dB_spec[idx_2kHz][x] > 10 and dB_spec[idx_5kHz][x] > 10 and dB_spec[idx_3kHz][x] > 10:
            if t[x] > 120:
                pass
            else:
                di_tag = t[x]

    return t[-1], di_tag + 0.01


def get_files(file_path):
    with os.scandir(file_path) as it:
        for entry in it:
            if entry.is_file() and entry.name.endswith('.pcm'):
                info = (entry.path, Path(entry.path).parts[-2])
                global file_list
                file_list.append(info)
            elif entry.is_dir():
                get_files(entry.path)
    return file_list


if __name__ == '__main__':
    if len(sys.argv) != 4:
        print("Usage: python3 audio_tag_split.py [audio_path] [total_channel] [people_channel]")
        sys.exit(1)

    # 获取音频路径
    audio_path = sys.argv[1]
    # 输入声道数
    total_channel = int(sys.argv[2])
    # 输入人声声道数
    people_channel = int(sys.argv[3])

    length, tag = find_voice_tag(audio_path, total_channel, people_channel)
    print(f'音频：{audio_path}, 总时长：{length}, 滴声：{tag}')
    process_file = audio_path.replace(".pcm", "_down.pcm")
    cmd = f'ffmpeg -ac {total_channel} -f s16le -ar 16000 -acodec pcm_s16le -i {audio_path} -ss {tag} -t {length-10} -f s16le -ar 16000 -ac {total_channel} {process_file}'
    os.system(cmd)
    print('切割音频成功：', process_file)

